package cn.johnyu.service;

import cn.johnyu.util.MilvusUtil;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResultData;
import io.milvus.grpc.SearchResults;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.IntStream;

public class MilvusService {
    protected static final Logger logger = LoggerFactory.getLogger(MilvusService.class);

    // 检查集合是否存在的逻辑
    public static boolean isCollectionExists(MilvusServiceClient milvusServiceClient, String collectionName, boolean close) {
        logger.info("检查集合" + collectionName + "是否存在...");
        HasCollectionParam hasParam = HasCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();

        R<Boolean> response = milvusServiceClient.hasCollection(hasParam);
        if (close) {
            milvusServiceClient.close();
        }
        logger.info(response.getData() ? "集合已经存在" : "集合404");
        return response.getData(); // 示例返回值
    }

    /**
     * 创建Collection
     * 三个字段： id、text、embedding
     */
    public static int createCollection(MilvusServiceClient milvusServiceClient, String collectionName, int dim) {

        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true) // 主键(必须指定)
                .withAutoID(true)  // 让 Milvus 自动生成主键 ID（可选）
                .build();

        FieldType textField = FieldType.newBuilder()
                .withName("text")
                .withDescription("原文")
                .withDataType(DataType.VarChar)
                .withMaxLength(512)
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("embedding")
                .withDescription("向量字段")
                .withDataType(DataType.FloatVector)
                .withDimension(dim)
                .build();

        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .addFieldType(idField)
                .addFieldType(textField)
                .addFieldType(vectorField)
                .withCollectionName(collectionName)
                .withShardsNum(2) // 分片数
                .build();
        //创建集合返回一个状态包装对象
        logger.info("正在创建集合..." + collectionName);
        R<RpcStatus> rpcStatusR = milvusServiceClient.createCollection(createCollectionParam);
        milvusServiceClient.flush(FlushParam.newBuilder()
                .withCollectionNames(Collections.singletonList(collectionName))
                .build());
        logger.info(rpcStatusR.getStatus() == 0 ? "创建集合成功" : "创建集合失败");
        return rpcStatusR.getStatus();
    }

    /**
     * 创建索引
     *
     * @param fieldName: 指定需要索引的字段
     * @return 0 成功，其它 失败
     */
    public static int createIndex(MilvusServiceClient milvusServiceClient, String collectionName, String fieldName) {
        logger.info("正在为集合创建索引..." + collectionName + " : " + fieldName);
        CreateIndexParam indexParam = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName(fieldName)
//                .withIndexName("basic_index")
                .withIndexType(IndexType.FLAT) //暴力全量对比
                .withMetricType(MetricType.L2) // 计算向量之间的距离(欧氏距离)
                .withExtraParam("{}") // 额外参数（可选）,因采用FLAT索引，所以不需要额外参数
                .build();
        R<RpcStatus> rpcStatusR = milvusServiceClient.createIndex(indexParam);
        logger.info(rpcStatusR.getStatus() == 0 ? "索引创建集合成功" : "索引创建集合失败");
        return rpcStatusR.getStatus(); // 示例返回值
    }

    /**
     * @param texts：  原文的列表
     * @param vectors 与原文对应的向量的列表
     * @return
     */
    public static long insertData(MilvusServiceClient milvusServiceClient, String collectionName, List<String> texts, List<List<Float>> vectors) {
        logger.info("正在插入数据..." + collectionName + "....共有 " + texts.size() + " 条数据。");

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(Arrays.asList(
                        new InsertParam.Field("text", texts),
                        new InsertParam.Field("embedding", vectors)
                ))
                .build();
        R<MutationResult> insertResultWrapper = milvusServiceClient.insert(insertParam);
        logger.info("成功插入数据的行数: " + insertResultWrapper.getData().getInsertCnt());
        //  刷新
        milvusServiceClient.flush(FlushParam.newBuilder()
                .withCollectionNames(Collections.singletonList(collectionName))
                .build());

        // 加载集合
        milvusServiceClient.loadCollection(LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build());
        // 等一会（或检测状态）
        // Thread.sleep(1000);
        return insertResultWrapper.getData().getInsertCnt(); // 示例返回值
    }

    /**
     * @param vector : 查询的向量
     * @return 列表中的项目为包含 “ID，相似度，原文” 的Map
     */
    public static List<Map<String, String>> searchData(MilvusServiceClient milvusServiceClient, String collectionName, List<Float> vector) {

        logger.info("正在搜索数据..." + vector.stream()
                        .map(String::valueOf)
                        .limit(3)
                .reduce("", (a, b) -> a + "," + b));

        List<Map<String, String>> rs = new ArrayList<>();

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withMetricType(MetricType.L2)
                .withOutFields(Arrays.asList("id", "text"))
                .withTopK(3) // 返回前3个最相似的向量
                .withVectors(Collections.singletonList(vector)) //
                .withVectorFieldName("embedding")
                .withParams("{\"nprobe\": 10}")
                .build();
        // 执行搜索，但返回的结果比较复杂，需要借助于SearchResultsWrapper进行封装处理
        R<SearchResults> searchResultsR = milvusServiceClient.search(searchParam);
        SearchResultData searchResultData = searchResultsR.getData().getResults();

        SearchResultsWrapper wrapper = new SearchResultsWrapper(searchResultData);
        List<String> texts = (List<String>) wrapper.getFieldWrapper("text").getFieldData();

        // 假设你只搜索了一个向量（如 vectors.get(0)）
        List<SearchResultsWrapper.IDScore> results = wrapper.getIDScore(0);
        for (int i = 0; i < results.size(); i++) {
            Map<String, String> map = new HashMap<>();
            map.put("id", results.get(i).getLongID() + "");
            map.put("score", results.get(i).getScore() + "");
            map.put("text", texts.get(i));
            rs.add(map);
            logger.info("欧氏距离: " + map.get("score")+",原文: " + map.get("text"));
        }
        return rs;
    }

    public static void main(String[] args) {
        int dim = 1024;
        MilvusServiceClient milvusServiceClient = MilvusUtil.getClient("localhost", 19530);
        String collectionName = "test_collection";
        boolean exists = MilvusService.isCollectionExists(milvusServiceClient, collectionName, false);
        if (!exists) {
            MilvusService.createCollection(milvusServiceClient, collectionName, dim);
            MilvusService.createIndex(milvusServiceClient, collectionName, "embedding");
        }

        // 模拟文本数据
        List<List<Float>> vectors=mockVectors(dim,3);
        List<String> texts=mockTexts(3);

        long count = MilvusService.insertData(milvusServiceClient, collectionName, texts, vectors);
        searchData(milvusServiceClient, collectionName, vectors.get(0));

        MilvusUtil.close();
    }

    private static List<String> mockTexts(int count) {
        List<String> texts = IntStream.range(0, count).mapToObj(i -> "中华人民共和国刑法 " + i)
                .collect(Collectors.toList());
        return  texts;
    }
    private static List<List<Float>> mockVectors(int dim,int count) {
        List<List<Float>> vectors = IntStream.range(0, count)
                .mapToObj(i -> new Random()
                        .doubles(dim)                  // 生成 double 流
                        .mapToObj(d -> (float) d)      // 转为 Float 类型
                        .collect(Collectors.toList())) // 收集成 List<Float>
                .collect(Collectors.toList());
        return vectors;
    }
}
